# encoding: utf-8
# @Time:    :2025/2/6 21:39

from torch.utils.data import Dataset


class NextTokenDataset(Dataset):
    def __init__(self, data_path: str, split_token: str = "__"):
        """

        :param data_path:
        """
        self.split_token = split_token
        self.data = []
        self.label = []
        self.load_data(data_path)

    def load_data(self, data_path):
        with open(data_path, "r", encoding="utf-8") as f:
            for line in f:
                text, label = line.strip().split(self.split_token)
                self.data.append(text)
                self.label.append(int(label))

    def __len__(self):
        return len(self.data)

    def __getitem__(self, idx):
        return self.data[idx], self.label[idx]


if __name__ == "__main__":
    pass
